% %clear;
% clc;
% close all;

% Steady State Values
function ER = Simulate_Concentration_HighOLNSDF_func(N,o)

A= 1;
W = 324;
P = 66;
% N=8;
%Model Parameters
h=0.5;
delta = 0.055;
phi = 0.05;
gamma = 5;
psik=0.78;
eta=0.525;
alph=0.99;
sigmaes =2.5;
nu=5;
% o=5.75;
theta=(1-gamma)/(1-eta);
sigmaf=(sigmaes-1)/sigmaes;
rho_w = 0.95;
rho_a = 0.88;
rho_p=0.965;
rho_x=0;
sd_w =0.025;
sd_p=0.0285;
sd_a=0.005;
cov_wp =0.011;
cov_wa=0.000231;
cov_pa=-0.00024;
wbar=log(W);
pibar=log(P);
abar=log(A);
rng(7,'twister');

NumSim = 1000;
w=zeros(59,NumSim);
pi=zeros(59,NumSim);
c=zeros(59,NumSim);
a=zeros(59,NumSim);
WS=zeros(59,NumSim);
PS=zeros(59,NumSim);
AS=zeros(59,NumSim);
u=zeros(59,NumSim);
logu=zeros(59,NumSim);
pSstar=zeros(59,NumSim);
kstar=zeros(59,NumSim);
ktilde=zeros(59,NumSim);
sstar=zeros(59,NumSim);
dstar=zeros(59,NumSim);
zdstar=zeros(59,NumSim);
ret=zeros(59,NumSim);
retss=zeros(59,NumSim);


for j= 1:NumSim

    % Disturbances in log output and log productivity shock,
    % eps = (eps_X, eps_theta) are bivariate normal with zero mean

    mu_eps = [0 0 0];
    sigma_eps = [sd_w^2 0 0; 0 sd_p^2 0; 0 0 sd_a^2]; % covariance matrix


    eps = mvnrnd(mu_eps,sigma_eps,59);
    w(1,j) =  wbar + eps(1,1);
    pi(1,j) =  pibar + eps(1,2);
    c(1,j)=w(1,j)-pi(1,j);
    a(1,j)=abar+eps(1,3);
    for i = 1:58
        w(i+1,j) =  rho_w*w(i,j) + eps(i+1,1);
        pi(i+1,j) =  rho_p*pi(i,j) + eps(i+1,2);
        c(i+1,j)=w(i+1,j)-pi(i+1,j);
        a(i+1,j)=rho_a*a(i,j)+eps(i+1,3);
        WS(i+1,j)=exp(w(i+1,j));
        PS(i+1,j)=exp(pi(i+1,j));
        AS(i+1,j)=exp(a(i+1,j));
    end

end
syms z 
x1=exp(z)/(1+exp(z));
x0=log(1+exp(z))-x1*z;
k_0=(log(alph)+x0)/(1-x1);
k_w=(rho_w-1)*(eta-1)/(x1*rho_w-1);
k_pi=(rho_p-1)*(theta*(eta-1)+1)/(theta*(x1*rho_p-1));
eqn1=z-(k_0+k_w*wbar+k_pi*pibar)==0;
S=vpasolve(eqn1,z,0.5);
lc=double(S);
disp(lc)
varkappa1 = exp(lc)/(1+exp(lc));
varkappa0=log(1+exp(lc))-varkappa1*lc;
kappa0=(log(alph)+varkappa0)/(1-varkappa1);
kappa_w=((rho_w-1)*(eta-1)/(varkappa1*rho_w-1));
kappa_p=((rho_p-1)*(theta*(1-eta)+1))/(theta*(varkappa1*rho_p-1));
B0=theta*log(alph);
B_w=(rho_w-1)*(theta*(1-eta)-1)+(theta-1)*kappa_w*(varkappa1*rho_w-1);
B_p=(rho_p-1)*(theta*(eta-1)+1)+(theta-1)*kappa_p*(varkappa1*rho_p-1);
b_w=-eta*theta+(theta-1)*(varkappa1*kappa_w+1);
b_p=eta*theta-(theta-1)*(1-varkappa1*kappa_p);
%Calculation of Punishment Steady State
v=(sigmaes*(1-psik)+psik)/sigmaes;
M=W^(1/sigmaes)*P^sigmaf;
% fun=@(k) Cap(k,N);
options = optimoptions('fsolve','TolFun',1e-20,'StepTolerance',1e-20,'Display','iter','MaxFunctionEvaluations',15000,'MaxIterations',200000); % Option to display output
[Ktilde,fval,exitflag]=fsolve(@(k) Cap(k,N),15,options);
ptilde=M*phi*(A)^(-1/sigmaes)*(Ktilde)^(-psik/sigmaes)*N^(-1/sigmaes);
Upstilde=M*phi*(A)^sigmaf*(Ktilde)^(psik*sigmaf)*N^(-1/sigmaes);
UpsKtilde=M*phi*(A)^(sigmaf)*(Ktilde)^(-v)*N^(-(sigmaes+1)/sigmaes)*((N*sigmaes-1)/sigmaes);
Dtilde=(ptilde-h)*A*(Ktilde)^psik-delta*Ktilde-o*Ktilde;
Stilde=(alph/(1-alph))*Dtilde;
T0Ktilde=nu;
T2Ktilde=alph*nu*delta*(2-delta);
T1Ktilde=-(nu*(1+alph*delta*(2-delta))+alph*psik*(UpsKtilde*v+(psik-1)*h*A*Ktilde^(psik-1)));
TOMKtilde=alph*(psik*(UpsKtilde-h*A*Ktilde^(psik-1))-o+(1-delta));
Zk1tilde=(2*T2Ktilde)^(-1)*(-T1Ktilde+(T1Ktilde^2-4*T2Ktilde*T0Ktilde)^0.5);
Zk2tilde=(2*T2Ktilde)^(-1)*(-T1Ktilde-(T1Ktilde^2-4*T2Ktilde*T0Ktilde)^0.5);
Zktilde=min(Zk1tilde,Zk2tilde);
Zkwtilde=-((psik*UpsKtilde/sigmaes)+TOMKtilde*B_w)/(T2Ktilde*(rho_w+Zktilde)+T1Ktilde);
Zkptilde=-((psik*(sigmaes-1)*UpsKtilde/sigmaes)+TOMKtilde*(B_p+(1-rho_p)))/(T2Ktilde*(rho_p+Zktilde)+T1Ktilde);
Zkatilde=-((psik*(sigmaes-1)*UpsKtilde/sigmaes))/(T2Ktilde*(rho_a+Zktilde)+T1Ktilde);
Zdwtilde=(Upstilde*(1/sigmaes)-Ktilde*Zkwtilde)/Dtilde;
Zdptilde=(Upstilde*((sigmaes-1)/sigmaes)-Ktilde*Zkwtilde)/Dtilde;
Zdatilde=(Upstilde*A^((sigmaes-1)/sigmaes)*((sigmaes-1)/sigmaes)-h*A*(Ktilde)^psik-Ktilde*Zkatilde)/Dtilde;
Zdktilde=(Ktilde/Dtilde)*((1-delta)-Zktilde-o);
Zsktilde=(1-alph)*Zdktilde*Zktilde/(1-alph*Zktilde);
Zswtilde=(B_w+(1-alph)*Zdwtilde*rho_w+Zkwtilde*(alph*Zsktilde+(1-alph)*Zdktilde))/(1-alph*rho_w);
Zsptilde=(B_p+(1-rho_p)+(1-alph)*Zdptilde*rho_p+Zkptilde*(alph*Zsktilde+(1-alph)*Zdktilde))/(1-alph*rho_p);
Zsatilde=((1-alph)*Zdatilde*rho_a+Zkatilde*(alph*Zsktilde+(1-alph)*Zdktilde))/(1-alph*rho_a);
Zk0tilde=log(Ktilde)*(T0Ktilde+T1Ktilde+T2Ktilde)/(T2Ktilde+(1+Zktilde)+T1Ktilde);
Zd0tilde=log(Dtilde)-(1/Dtilde)*(wbar*Upstilde/sigmaes+abar*(sigmaf-h*A*Ktilde)+pibar*sigmaf+log(Ktilde)*(sigmaf*psik-psik*h*A*Ktilde-(o+delta)*Ktilde+Zk0tilde));
Zs0tilde=(B0-alph+(1-alph)*(Zd0tilde-log(Dtilde))+(1-alph)*log(Stilde))/(1-alph);
%Calculation of u, K in Steady State Equilibrium
k=Ktilde;
y0=[0.1; 1];
options = optimoptions('fsolve','TolFun',1e-20,'StepTolerance',1e-20,'Display','iter','MaxFunctionEvaluations',15000,'MaxIterations',200000); % Option to display output
[g,fval,exitflag]=fsolve(@(y)IC(y,Ktilde,N),y0,options);
ustaro=g(1);
Kstaro=g(2);
ustar=real(ustaro);
Kstar=real(Kstaro);
pstar=M*phi*(A)^(-1/sigmaes)*(ustar*Kstar)^(-psik/sigmaes)*N^(-1/sigmaes);
Icoststar=delta*Kstar;
Dstar=(pstar-h)*A*(ustar*Kstar)^psik-Icoststar-o*Kstar;
Sstar=(alph*Dstar)/(1-alph);
Diff=(pstar-h)*A*Kstar^psik*(1-ustar^psik)-Sstar+Stilde;
%Equilbrim Path Analysis
Upsstar=M*phi*A^sigmaf*Kstar^(psik*sigmaf)*N^(-1/sigmaes)*ustar^(-psik/sigmaes);
Upsustar=Upsstar*ustar^psik;
UpsKstar=psik*sigmaf*Upsustar/Kstar;
Hu0=psik*(h*A*(ustar*Kstar)^psik-(Upsstar/sigmaes)*(1+ustar^psik*(sigmaes-1)));
Ostar=h*(N*A)^(1/sigmaes)*Kstar^(psik/sigmaes)*ustar^(psik*(sigmaes+1)/sigmaes)/(M*phi);
Xistar=-(sigmaes-1)*ustar^psik+sigmaes*Ostar;
Hua=Upsstar*sigmaf-h*A*Kstar^psik;
HuK=Hua*psik;
Hud=Upsustar*sigmaf-h*A*(ustar*Kstar)^psik;
TK2Kstar=alph*nu*(delta+(1-Xistar)*(1-delta));
TMKstar=alph*((UpsKstar-psik*h*A*ustar^psik*Kstar^(psik-1))*(1-Xistar*ustar^(-(psik*(sigmaes-1)/sigmaes)))-psik*Xistar*h*A*(ustar^(psik/sigmaes)-1)*Kstar^(psik-1)+(1-delta-o)*(1-Xistar));
TK0Kstar=nu*(1-Xistar)+psik*Ostar*Xistar*(1+alph*((UpsKstar-psik*h*A*ustar^psik*Kstar^(psik-1))*ustar^(-(psik*(sigmaes-1)/sigmaes))- psik*Xistar*h*A*(ustar^(psik/sigmaes)-1)*Kstar^(psik-1)-(1-delta-o)));
TK1Kstar=-nu*(1-Xistar)-alph*(UpsKstar*v*(1-Xistar*ustar^(-(psik*(sigmaes-1)/sigmaes)))+(psik-1)*psik*Xistar*h*A*(ustar^(psik/sigmaes)-1)*Kstar^(psik-1)+nu*(delta*(1-Xistar)-(1-delta)*((1-Xistar)-delta*Xistar*psik*Ostar)));
Tu1Kstar=alph*(psik*UpsKstar*sigmaf-psik*h*A*Kstar^(psik-1)*(psik+Xistar*ustar^(-(psik*(sigmaes-1)/sigmaes))*(psik*sigmaf*ustar^psik-psik)+Xistar*psik/sigmaes)+(1-delta)*nu*delta*Xistar*(Ostar*(sigmaes+1)-(sigmaes-1)*ustar^psik));
Tu0Kstar=(Ostar*(sigmaes+1)-(sigmaes-1)*ustar^psik)*Xistar*(nu*delta*psik+alph*((UpsKstar-psik*h*A*ustar^psik*Kstar^(psik-1))*ustar^(-(psik*(sigmaes-1)/sigmaes))-psik*h*Xistar*A*(ustar^(psik/sigmaes)-1)*Kstar^(psik-1)+(o-(1-delta))));
TW0Kstar=-Xistar*Ostar*(1+alph*((UpsKstar-psik*h*A*ustar^psik*Kstar^(psik-1))*ustar^(-(psik*(sigmaes-1)/sigmaes))-psik*h*Xistar*A*(ustar^(psik/sigmaes)-1)*Kstar^(psik-1)-(1-delta-o)));
TW1Kstar=alph*(UpsKstar/sigmaes*(1-Xistar*ustar^(-(psik*(sigmaes-1)/sigmaes)))-(1-delta)*nu*delta*Xistar*Ostar);
TP0Kstar=(sigmaes-1)*TW0Kstar;
TP1Kstar=(sigmaes-1)*TW1Kstar;
TA0Kstar=-TW0Kstar;
TA1Kstar=alph*((sigmaf*UpsKstar-psik*h*A*ustar^(psik/sigmaes)*(Kstar)^(psik-1))*(1-Xistar*ustar^(-(psik*(sigmaes-1)/sigmaes)))-psik*h*Xistar*A*(ustar^(psik/sigmaes)-1)*Kstar^(psik-1)+(1-delta)*nu*delta*Xistar*Ostar);
TK2TildeStar=-nu*delta*Xistar;
%Solving for optimal capital
%Define auxiliary variables
Hdk=(1/Dstar)*(1-delta-o)*Kstar+psik*Hud*(Hu0-(1-ustar^psik)*HuK)/(Dstar*Hu0);
Hdw=(1/(sigmaes*Dstar*Hu0))*(Hu0*Upsustar-psik*(1-ustar^psik)*Upsstar*Hud);
Hdp=((sigmaes-1)/(sigmaes*Dstar*Hu0))*(Hu0*Upsustar-psik*(1-ustar^psik)*Upsstar*Hud);
Hda=(Hud/(Dstar*Hu0))*(Hu0-psik*(1-ustar^psik)*Hua);
Hds=(psik*Hud)*Sstar/(Dstar*Hu0);
Hdstildestar=-(psik*Hud)*Stilde/(Dstar*Hu0);
CXi00Kstar=TK0Kstar+(Tu0Kstar/Hu0)*(Sstar*(1-alph)*Hdk+Stilde*Zsktilde*((1-alph)*Hdstildestar-1)-(1-ustar^psik)*HuK);
CXi01Kstar=-alph*TK0Kstar-(Tu0Kstar/Hu0)*(Sstar*Kstar*(1-alph)-alph*(Stilde*Zsktilde+(1-ustar^psik)*HuK));
CXi10Kstar=TK1Kstar+(Tu1Kstar/Hu0)*(Sstar*(1-alph)*Hdk+Stilde*Zsktilde*((1-alph)*Hdstildestar-1)-(1-ustar^psik)*HuK);
CXi11Kstar=-alph*TK1Kstar-(Tu1Kstar/Hu0)*(Sstar*Kstar*(1-alph)-alph*(Stilde*Zsktilde+(1-ustar^psik)*HuK));
That0K=real(CXi00Kstar);
That1K=real(CXi01Kstar+CXi10Kstar);
That2K=real(TK2Kstar+alph*CXi11Kstar);
That3K=-alph*real(TK2Kstar);
q=[That3K That2K That1K That0K];
r=roots(q);
rreal=r(imag(r)==0);
rmin=min(rreal);
Zkkstar=rmin;
Zdsstar=Hds;
Zdstildestar=Hdstildestar;
Zskstar=(1-alph)*(Hdk-Kstar*(1-delta-o)*Zkkstar+Zdstildestar*Zsktilde)/(1-alph*Zkkstar);
Zdkstar=Hdk-Kstar*(1-delta-o)*Zkkstar;
syms x y
eqn1 =TP0Kstar+rho_p*(TP1Kstar+TMKstar*(B_p+(1-rho_p)))+(1/Hu0)*(Tu0Kstar+rho_p*Tu1Kstar)*(Sstar*y-Stilde*Zsptilde-(1-ustar^psik)*Upsstar*sigmaf)+...
    (TK1Kstar+rho_p*TK2Kstar)*x==0;
eqn2= ((1-alph)*rho_p*(Hdp-Kstar*(1-delta-o)*x+Zdstildestar*Zsptilde)+x*((1-alph)*(Hdk-Kstar*(1-delta-o)*Zdkstar)+alph*Zkkstar)+B_p+(1-rho_p))/((1-alph*rho_p)+(alph-1)*rho_p*Zdsstar)==y;
sol=solve([eqn1,eqn2],[x,y]);
xSol=sol.x;
ySol=sol.y;
Zkpstar=double(xSol);
Zspstar=double(ySol);
Zdpstar=Hdp-Kstar*(1-delta-o)*Zkpstar;
syms m n
eqn1 =TW0Kstar+rho_w*(TW1Kstar+TMKstar*B_w)+(1/Hu0)*(Tu0Kstar+rho_w*Tu1Kstar)*(Sstar*n-Stilde*Zswtilde-(1-ustar^psik)*Upsstar/sigmaes)+...
    (TK1Kstar+rho_w*TK2Kstar)*m==0;
eqn2= ((1-alph)*rho_w*(Hdw-Kstar*(1-delta-o)*m+Zdstildestar*Zswtilde)+m*((1-alph)*(Hdk-Kstar*(1-delta-o)*Zdkstar)+alph*Zkkstar)+B_w)/((1-alph*rho_w)+(alph-1)*rho_w*Zdsstar)==n;
sol=solve([eqn1,eqn2],[m,n]);
mSol=sol.m;
nSol=sol.n;
Zkwstar=double(mSol);
Zk0star=-(log(ustar)*(Tu0Kstar+Tu1Kstar)+log(Kstar)*(TK0Kstar+TK1Kstar+TK2Kstar)+wbar*(TW0Kstar+TW1Kstar)+pibar*(TP0Kstar+TP1Kstar)+abar*(TA0Kstar+TA1Kstar)+alph*TMKstar+log(Ktilde)*TK2TildeStar)/(1-Tu1Kstar*HuK);
Zswstar=double(nSol);
Zdwstar=Hdw-Kstar*(1-delta-o)*Zkwstar;
syms m n
eqn1 =TW0Kstar+ rho_a*TW1Kstar+(1/Hu0)*(Tu0Kstar+rho_x*Tu1Kstar)*(Sstar*n-Stilde*Zsatilde-(1-ustar^psik)*Hua)+...
    (TK1Kstar+rho_x*TK2Kstar)*m==0;
eqn2= ((1-alph)*rho_a*(Hda-Kstar*(1-delta-o)*m+Zdstildestar*Zsatilde)+m*((1-alph)*(Hdk-Kstar*(1-delta-o)*Zdkstar)+alph*Zkkstar))/((1-alph*rho_a)+(alph-1)*rho_a*Zdsstar)==n;
sol=solve([eqn1,eqn2],[m,n]);
mSol=sol.m;
nSol=sol.n;
Zkastar=double(mSol);
Zsastar=double(nSol);
Zdastar=Hda-Kstar*(1-delta-o)*Zkastar;
syms m n
eqn1=((theta-1)*log(alph)+(1-alph)*(n-log(Dstar))+(1-alph)*log(Sstar))/(1-alph)==m;
eqn2=log(Dstar)-(1/Dstar)*(Upsustar*(wbar/sigmaes+sigmaf*pibar)+(abar+psik*log(ustar))*(Upsustar*sigmaf-h*A*(ustar*Kstar)^psik)+log(Kstar)*(psik*(Upsustar*sigmaf-h*A*(ustar*Kstar)^psik)-Kstar*(o+Zk0star-delta)))+...
    (1/Dstar)*(Upsustar*sigmaf-h*A*(ustar*Kstar)^psik)*(1/Hu0)*((Sstar*(m-log(Sstar))-Stilde*(Zs0tilde-log(Stilde))+(1-ustar^psik)*(Upsstar*(wbar/sigmaes+sigmaf*pibar)-abar*Hua-log(Kstar)*HuK)))==n;
sol=solve([eqn1,eqn2],[m,n]);
mSol=sol.m;
nSol=sol.n;
Zs0star=double(mSol);
Zd0star=double(nSol);
% Unconditional ERP Calculation
Lw=((1-alph)/alph)*(1/Dstar)*(Zdwstar+Zdstildestar*Zswtilde+Zdsstar*Zswstar)+Zswstar;
Lp=((1-alph)/alph)*(1/Dstar)*(Zdpstar+Zdstildestar*Zsptilde+Zdsstar*Zspstar)+Zspstar;
La=((1-alph)/alph)*(1/Dstar)*(Zdastar+Zdstildestar*Zsatilde+Zdsstar*Zsastar)+Zsastar;
beta_w=-alph*b_w*Lw;
beta_p=-alph*(b_p-1)*Lp;
beta_wp=-alph*((b_p-1)*Lp+b_w*Lw);
beta_wa=-alph*b_w*La;
beta_pa=-alph*(b_p-1)*La;
ERP=beta_w*(sd_w)^2+beta_p*(sd_p)^2+beta_wp*cov_wp+beta_wa*cov_wa+beta_pa*cov_pa;
%Simulate ROI
for j=1:NumSim
    for i=1:59
        sstar(i,j)=Zs0star+kstar(i,j)*Zskstar+w(i,j)*Zswstar+pi(i,j)*Zspstar+a(i,j)*Zsastar;
        dstar(i,j)=Zd0star+kstar(i,j)*Zdkstar+w(i,j)*Zdwstar+pi(i,j)*Zdpstar+a(i,j)*Zdastar;
        zdstar(i,j)=sstar(i,j)-dstar(i,j);
    end
end
Mean_s=mean(mean(sstar(1:59)));
Mean_zd=mean(mean(zdstar(1:59)));
Mean_d=mean(mean(dstar(1:59)));
varkappd1=exp(Mean_zd)/(1+exp(Mean_zd));
varkappd0=log(1+exp(Mean_zd))-varkappd1*Mean_zd;
for j=1:NumSim
    for i=1:58
        ret(i+1,j)= varkappd0+varkappd1*zdstar(i+1,j)-zdstar(i,j)+dstar(i+1,j)-dstar(i,j);
    end
end
ER=mean(mean(ret(1:59)));

function C=Cap(k,N)
A= 1;
W = 324;
P = 66;
% N=8;
%Model Parameters
h=0.5;
delta = 0.055;
phi = 0.05;
gamma = 5;
psik=0.78;
eta=0.525;
alph=0.99;
sigmaes =2.5;
nu=5;
o=5.75;
sigmaf=(sigmaes-1)/sigmaes;
M=(W)^(1/sigmaes)*(P)^sigmaf;
v=(sigmaes*(1-psik)+psik)/sigmaes;
UpsKtilde=M*phi*(A)^sigmaf*(k)^(-v)*N^(-(sigmaes+1)/sigmaes)*((N*sigmaes-1)/sigmaes);
C=alph*(psik*(UpsKtilde-h*A*k^(psik-1))-o +(1-delta))-1;



function f= IC(y,k,N)

A= 1;
W = 324;
P = 66;
% N=8;
%Model Parameters
h=0.5;
delta = 0.055;
phi = 0.05;
gamma = 5;
psik=0.78;
eta=0.525;
alph=0.99;
sigmaes =2.5;
nu=5;
o=5.75;
sigmaf=(sigmaes-1)/sigmaes;
v=(sigmaes*(1-psik)+psik)/sigmaes;
r=y(2);
y=y(1);
M=(W)^(1/sigmaes)*(P)^((sigmaes-1)/sigmaes);
sigmaf=(sigmaes-1)/sigmaes;
ptilde=M*phi*(A)^(-1/sigmaes)*(k)^(-psik/sigmaes)*N^(-1/sigmaes);
Dtilde=(ptilde-h)*A*(k)^psik-(delta+o)*k;
pstar=M*phi*N^(-1/sigmaes)*(y*r)^(-psik/sigmaes)*A^(-1/sigmaes);
f(1)=Dtilde+(delta+o)*r-(pstar-h)*A*r^psik*(y^psik-(1-alph))/alph;
f(2)=-(1-alph*((1-delta)-o))*(pstar*(1-y^psik)+y^psik*sigmaes*(pstar-h))/pstar+alph*(psik*(sigmaf*W^(1/sigmaes)*P^sigmaf*phi*A^sigmaf*y^(psik*sigmaf)*r^(-v)*N^(-1/sigmaes)-h*A*y^psik*r^(psik-1))*(pstar*(1-y^(psik/sigmaes))+y^(psik/sigmaes)*sigmaes*(pstar-h))/pstar-psik*y^(psik*(sigmaes+1)/sigmaes)*(1-y^(psik/sigmaes))*h*A*r^(psik-1)*(pstar-sigmaes*(pstar-h))/pstar);





% tic